package com.feidee.fd.sml.algorithm.component.feature

import com.feidee.fdspark.transformer._
import org.apache.spark.ml.feature.{FFMNodesEncoder => FNE}
import org.apache.spark.ml.{Pipeline, PipelineModel}
import org.apache.spark.sql.DataFrame

/**
  * @Author songhaicheng
  * @Date 2019/2/18 19:36
  * @Description
  * @Reviewer
  */
case class FFMNodesEncoderParam(
                                 override val input_pt: String,
                                 override val output_pt: String,
                                 override val hive_table: String,
                                 override val flow_time: String,
                                 override val inputCol: String,
                                 override val outputCol: String,
                                 override val preserveCols: String,
                                 override val modelPath: String,
                                 // OneHot 编码列，多列用逗号分隔
                                 oneHotCols: String,
                                 // 离散值列，多列用逗号分隔
                                 vectorCols: String,
                                 // 连续变量列，多列用逗号分隔
                                 continuousCols: String,
                                 // 离散值列的词分隔符，默认空格
                                 wordDelimiter: String,
                                 // 是否处理 OneHot 和 离散值列出现的无效值（未出现在训练数据中），true 则给所有无效值分配一个 feature 值，false 会直接报错，默认 false
                                 handleInvalid: Boolean
                               ) extends FeatureParam {

  def this() = this(null, null, null, null, null, "features", null, null, null, null, null, " ", false)

  override def verify(): Unit = {
    require(!(!tool.isNotNull(oneHotCols) &&
      !tool.isNotNull(vectorCols) &&
      !tool.isNotNull(continuousCols)),
      "none will be encoded")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("oneHotCols" -> oneHotCols)
    map += ("vectorCols" -> vectorCols)
    map += ("continuousCols" -> continuousCols)
    map += ("wordDelimiter" -> wordDelimiter)
    map += ("handleInvalid" -> handleInvalid)
    map
  }
}


class FFMNodesEncoder extends AbstractFeatureEncoder[FFMNodesEncoderParam] {

  override def train(param: FFMNodesEncoderParam, data: DataFrame): PipelineModel = {
    val preserveCols = if (tool.isNotNull(param.preserveCols)) {
      param.preserveCols.split(",")
    } else {
      new Array[String](0)
    }
    val oCols = if (tool.isNotNull(param.oneHotCols)) {
      param.oneHotCols.split(",")
    } else {
      new Array[String](0)
    }
    val vCols = if (tool.isNotNull(param.vectorCols)) {
      param.vectorCols.split(",")
    } else {
      new Array[String](0)
    }
    val cCols = if (tool.isNotNull(param.continuousCols)) {
      param.continuousCols.split(",")
    } else {
      new Array[String](0)
    }
    val cols = preserveCols ++ Array[String](param.outputCol)

    val meta = new MetaStorage()
      .setModelType(ModelType.Feature)
      .setParameters(param.toMap)
      .setFields(data.schema.fieldNames)

    // 处理离散列，生成词组
    val spliters = new Array[WordSpliter](vCols.length)
    for (i <- vCols.indices) {
      val spliter = new WordSpliter()
          .setInputCol(vCols(i))
          .setOutputCol(tool.renameDuplicatedColName(s"${vCols(i)}_words", cols))
          .setDelimiter(param.wordDelimiter)
          .setToType(ToTypes.Array)
      spliters.update(i, spliter)
    }

    // 调用 FFMNodes 编码器
    val encoder = new FNE()
      .setOneHotCols(oCols)
      .setVectorCols(spliters.map(_.getOutputCol))
      .setContinuousCols(cCols)
      .setOutputCol(param.outputCol)
    if (param.handleInvalid) {
      encoder.setHandleInvalid("handle")
    }

    val colSelector = new ColSelector()
      .setCols(cols)

    val pipeline = new Pipeline().setStages(Array(meta) ++ spliters ++ Array(encoder, colSelector))
    val model = pipeline.fit(data)

    model
  }

}

object FFMNodesEncoder {
  def apply(paramStr: String): Unit = {
    new FFMNodesEncoder()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    FFMNodesEncoder(args(0))
  }
}
